{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c66c05bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tokenizers import Tokenizer\n",
    "from pathlib import Path\n",
    "import sys\n",
    "\n",
    "sys.path.append(str(Path(Path.cwd()).parent))\n",
    "\n",
    "from helper.sentence_destructor import destruct_sentence, mask_sentence, dissolve_word"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d52f2cc3",
   "metadata": {},
   "source": [
    "### Test dataset destruction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "558c577c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "مەن مەكتەپتىن قايتىپ <MASK> ئۇمۇ مەكتەپتىن <MASK> كەلدى، بىز ھەممىمىز مەكتەپتىن قايتىپ كەلدۇق!\n",
      "مەن مەكتەپتىن قايتىپ كەلدىم، ئۇمۇ مەكتەپتىن قايتىپ كەلدى،ب ىز ھەممىمىز مەكتەپتىن اقيتىپ كەلدۇق!\n",
      "----------------------------------------------------------------------------------------------------\n",
      "مەن مەكتەپتىن <MASK> كەلدىم، <MASK> مەكتەپتىن <MASK> كەلدى، بىز <MASK> مەكتەپتىن قايتىپ كەلدۇق!\n",
      "مەن مەكتگپتن قاىدپك گلمد، ئۇمۇ مەكەتپتن قايتىپ كەلدى، بى زھەمىممىر مەكتەپت ىن قايتخپ كەلدۇق!\n",
      "----------------------------------------------------------------------------------------------------\n",
      "مەن مەكتەپتىن <MASK> كەلدىم، <MASK> مەكتەپتىن قايتىپ كەلدى، بىز ھەممىمىز مەكتەپتىن قايتىپ كەلدۇق!\n",
      "مەن مەكتەپتن قاتيىپ كەلدېم،ئ ۇمۇ مەكتەپتىن قايتىپكەلدى، بىز ھەمممىز مەكتەپتن قايتىپ كەلدۇق!\n",
      "----------------------------------------------------------------------------------------------------\n",
      "مەن مەكتەپتىن قايتىپ <MASK> <MASK> مەكتەپتىن قايتىپ <MASK> بىز ھەممىمىز مەكتەپتىن قايتىپ كەلدۇق!\n",
      "مگن مەكدەپتىن قايتىپكەلدىم، ئۇمۇ مگكتەپتىن قايدپ كگلدى، بخز ھەممىمز مەكتەپدىن قاىتىپ كەلدۇق!\n",
      "----------------------------------------------------------------------------------------------------\n",
      "مەن مەكتەپتىن قايتىپ كەلدىم، ئۇمۇ مەكتەپتىن قايتىپ <MASK> بىز ھەممىمىز مەكتەپتىن قايتىپ كەلدۇق!\n",
      "مەن مەكتەپتىن قايدىپ كەلدم، ئۇمۇ مەكتاپتىن قايتىپ كەلدى، بىزھگمممز مەكتەپتىن قايتىپ كەلدۇق!\n",
      "----------------------------------------------------------------------------------------------------\n",
      "مەن مەكتەپتىن قايتىپ كەلدىم، ئۇمۇ مەكتەپتىن قايتىپ كەلدى، بىز ھەممىمىز مەكتەپتىن قايتىپ كەلدۇق!\n"
     ]
    }
   ],
   "source": [
    "\n",
    "tokenizer: Tokenizer = Tokenizer.from_file(\"../../config/input_tokenizer.json\")\n",
    "\n",
    "text = \"مەن مەكتەپتىن قايتىپ كەلدىم، ئۇمۇ مەكتەپتىن قايتىپ كەلدى، بىز ھەممىمىز مەكتەپتىن قايتىپ كەلدۇق!\"\n",
    "\n",
    "for i in range(5):\n",
    "    masked_text = mask_sentence(text, 0.2)\n",
    "    destructed = destruct_sentence(tokenizer, text, 0.15)\n",
    "\n",
    "    print(masked_text)\n",
    "    print(destructed)\n",
    "    print(\"-\"*100)\n",
    "print(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "107144b9",
   "metadata": {},
   "source": [
    "### Read tsv files and write them to lmdb."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "995f10c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "file: sub_sentences_39.tsv, id: 2000:   0%|                                 | 0/437 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "file: sub_sentences_60.tsv, id: 4119000: 100%|████████████████████| 437/437 [06:19<00:00,  1.15it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Added 4119872 pairs to LMDB database.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "import csv\n",
    "import pickle\n",
    "from pathlib import Path\n",
    "from tokenizers import Tokenizer, decoders, pre_tokenizers\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "import time\n",
    "import lmdb\n",
    "\n",
    "input_tokenizer: Tokenizer = Tokenizer.from_file(str(Path.cwd().parent.parent / \"config/input_tokenizer.json\"))\n",
    "# output_tokenizer = input_tokenizer\n",
    "output_tokenizer: Tokenizer = Tokenizer.from_file(str(Path.cwd().parent.parent / \"config/output_tokenizer.json\"))\n",
    "output_tokenizer.pre_tokenizer = pre_tokenizers.Metaspace()\n",
    "output_tokenizer.decoder = decoders.Metaspace()\n",
    "\n",
    "#list all the tsv files in .data/text/by_sentences/*/tsv\n",
    "tsv_files: list[Path] = list((Path.cwd().parent.parent / \".data/text/by_sentences\").rglob(\"*.tsv\"))\n",
    "\n",
    "db = lmdb.Environment(str(Path.cwd().parent.parent / \".data/temp/by_sentences_lmdb\"), map_size=1024 * 1024 * 1024 * 15)\n",
    "\n",
    "id: int = 0\n",
    "\n",
    "tx = db.begin(write=True)\n",
    "def write_data(id: int, data: tuple[list[int], list[int]]):\n",
    "    global db\n",
    "    global tx\n",
    "    tx.put(f\"{id}\".encode(), pickle.dumps(data))\n",
    "    if id % 1000 == 0:\n",
    "        tx.commit()\n",
    "        tx = db.begin(write=True)\n",
    "\n",
    "progress_bar = tqdm(tsv_files, desc=\"Processing tsv files\", ncols=100)\n",
    "for tsv_file in progress_bar:\n",
    "    with open(tsv_file, \"r\") as file:\n",
    "        reader = csv.reader(file, delimiter=\"\\t\")\n",
    "        header = next(reader) #skip the header\n",
    "        #find the \"title_ug\" index\n",
    "        title_ug_index = header.index(\"title_ug\")\n",
    "        for row in reader:\n",
    "            target_text = row[title_ug_index]\n",
    "            target_ids: list[int] = output_tokenizer.encode(f\"<SOS>{target_text}<EOS>\").ids\n",
    "\n",
    "            masked_inputs = mask_sentence(target_text, 0.15)\n",
    "            dissolved_sentence = dissolve_word(target_text, 0.07)\n",
    "            \n",
    "            # misspelled_inputs = [destruct_sentence(input_tokenizer, target_text, random.uniform(0.02, 0.2)) for _ in range(3)]\n",
    "            # misspelled_inputs =  misspelled_inputs + [dissolved_sentence, masked_inputs]\n",
    "            # misspelled_inputs =  [dissolved_sentence, masked_inputs]\n",
    "            # misspelled_inputs =  [masked_text]\n",
    "            misspelled_inputs = [\"\"]\n",
    "            for misspelled_input in misspelled_inputs:\n",
    "                misspelled_input_ids: list[int] = input_tokenizer.encode(f\"<SOS>{misspelled_input}<EOS>\").ids\n",
    "                \n",
    "                if len(misspelled_input_ids) > 512:\n",
    "                    continue\n",
    "                pair = (misspelled_input_ids, target_ids)\n",
    "                write_data(id, pair)\n",
    "                id += 1\n",
    "            \n",
    "            #decrease the progress bar update frequency\n",
    "                if id % 1000 == 0:\n",
    "                    progress_bar.set_description(f\"file: {tsv_file.name}, id: {id}\")\n",
    "tx.put(b\"count\", pickle.dumps(id))\n",
    "tx.commit()\n",
    "\n",
    "print(f\"Added {id} pairs to LMDB database.\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "spell-correction",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.2"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
